import torch.nn as nn
import math
import numpy as np
import torch
import datetime

def dataparallel(model, ngpus, gpu0=0):  # data parallel
    if ngpus == 0:
        assert False, "only support gpu mode"
    gpu_list = list(range(gpu0, gpu0 + ngpus))
    assert torch.cuda.device_count() >= gpu0 + ngpus, "Invalid Number of GPUs"
    if isinstance(model, list):
        for i in range(len(model)):
            if ngpus >= 2:
                if not isinstance(model[i], nn.DataParallel):
                    model[i] = torch.nn.DataParallel(model[i], gpu_list).cuda()
            else:
                model[i] = model[i].cuda()
    else:
        if ngpus >= 2:
            if not isinstance(model, nn.DataParallel):
                model = torch.nn.DataParallel(model, gpu_list).cuda()
        else:
            model = model.cuda()
    return model


def writelog(log_str, log_file):

    fo = open(log_file, "a")
    fo.writelines(log_str)
    fo.close()


def draw_result(log_file):  # 不需要用

    train_loss = []
    threshold = []
    train_F1 = []
    test_loss = []
    P = []
    R = []
    F1 = []
    count = 0
    with open(log_file, 'r') as f:
        file_list = f.readlines()
        for file in file_list:
            file = file.strip('\n')
            results = file.split('\t')
            print(count)
            if count % 2 == 0:
                trainloss = results[1].split(':')[1]
                thres = results[2].split(':')[1]
                trainF1 = results[3].split(':')[1]
                train_loss.append(trainloss)
                threshold.append(thres)
                train_F1.append(trainF1)
            else:
                testloss = results[1].split(': ')[1]
                p = results[3].split(': ')[1]
                r = results[4].split(': ')[1]
                f1 = results[5].split(': ')[1]
                test_loss.append(testloss)
                P.append(p)
                R.append(r)
                F1.append(f1)
            count += 1
    print(count)
    y = range(1, count / 2 + 1)

    
    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    #ax1.plot(train_loss, label = "train_loss", color = 'r')
    ax1.plot(test_loss, label = "test_loss", color = 'r')
    ax1.set_ylabel("Loss")
    ax1.set_xlabel("Epoch")

    ax2 = ax1.twinx()
    ax2.plot(F1, label = "F1", color = 'b')
    #ax2.plot(train_F1, label = "train_F1", color = 'b')
    ax2.set_ylabel("F1")
    ax1.legend(loc='upper right')
    ax2.legend(loc='upper left')
    ax2.set_title("Test_loss and Test_F1")
    

    # fig.savefig("/home/zhengyihan/zero-shot-learning/model/loss_F1.png")
    """

    pl.plot(y, train_loss, 'r')
    pl.plot(y, test_loss, 'g')
    

    pl.title('Plot of training loss vs. testing loss')  # give plot a title
    pl.xlabel('Epoch')  # make axis labels
    pl.ylabel('Loss')
    pl.savefig("/home/zhengyihan/zero-shot-learning/model/ML_ZSL/loss.png")
    """

